import matplotlib.pyplot as plt
import cv2
import numpy as np
from skimage.transform import resize
from skimage.filters import *
from skimage import exposure
import pandas as pd

import numpy as np
from skimage.feature import greycomatrix, greycoprops
from skimage import io, color, img_as_ubyte
import os

# Parameters
resize_resolution = (1080, 1640)
gaussian_blur_kernel = (3, 3)
gamma_adjust = 0.8

dilation_iters = 4
dilation_structuring_element = (10, 3)

min_area = 200
cntr = 0

def preprocess_image(img):
    """
        Image Preprocessing - Image Resize, Apply Gaussian Blur, Adjust contrast
    """
    # Perform Image resizing
    img = cv2.resize(img, resize_resolution, interpolation = cv2.INTER_AREA)
    # Apply Gaussian Blur to smoothen image and remove noize
    img = cv2.GaussianBlur(img, gaussian_blur_kernel, 0)
    # Adjust contrast in the image
    img = exposure.adjust_gamma(img, gamma_adjust)
    return img

def local_thresholding(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = np.array(img)
    block_size = 35
    local_thresh = threshold_local(img, block_size, offset=10)
    binary_local = img > local_thresh
    img = np.array(binary_local).astype(np.uint8)
    return img
    
def remove_lines(img):
    # Horizontal lines
    lines = np.ones(img.shape)
    lines = lines.astype(np.uint8)
    
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40,2))
    detected_lines = cv2.morphologyEx(img, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
    cnts = cv2.findContours(detected_lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if len(cnts) == 2 else cnts[1]
    for c in cnts:
        cv2.drawContours(lines, [c], -1, 0, 8)
    img = cv2.subtract(lines, img)
    return img

def custom_dilation(img):
    """
        Perform morphological operation - Dilation with customized structuring element and cleanup the border pixels of the Image
    """
    kernel = np.ones((3, 3), np.uint8)
    img = cv2.erode(img, kernel)
    # Define custom rectangular structuring element
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, dilation_structuring_element)
    # Perform morphological dilation on the image for multiple iterations
    img = cv2.dilate(img, kernel, iterations=dilation_iters)
    # Clean the border pixels of the image
    img[0:25, :] = 0
    img[:, 0:25] = 0
    img[:, len(img[0])-25:len(img[0])-1] = 0
    img[len(img)-25:len(img)-1,: ] = 0
    return img

import os
from collections import Counter
from skimage import morphology
from glob import glob

for i, image in enumerate(glob('./images/*.jpg')):
    fig, (ax1, ax2, ax3, ax4, ax5, ax6) = plt.subplots(1, 6, figsize=(40, 60), dpi=200)

    input_image = cv2.imread(image)
    input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    ax1.imshow(input_image)
    ax1.title.set_text("Input Image")
    
    prep_image = preprocess_image(input_image)
    ax2.imshow(prep_image)
    ax2.title.set_text("Preprocessing")

    thresh = local_thresholding(prep_image)
    ax3.imshow(thresh, cmap="gray")
    ax3.title.set_text("Thresholding")
    
    thresh = remove_lines(thresh)
    ax4.imshow(thresh, cmap="gray")
    ax4.title.set_text("Line Removal")

    dilation = custom_dilation(thresh)
    ax5.imshow(dilation, cmap='gray')
    ax5.title.set_text("Dilation")
    
    max_area = (len(prep_image[0])-30) * (len(prep_image)-30)
    cnts, _ = cv2.findContours(dilation, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    bbox_image = prep_image.copy()
    
    cntr = 1
    for cntr in cnts:
        cntr += 1
        area = cv2.contourArea(cntr)
        # Small object Removal
        if area < min_area or area > max_area:
            continue
        # Bounding Box around word
        x, y, w, h = cv2.boundingRect(cntr)
        cv2.rectangle(bbox_image, (x-5, y-5), ((x + w) - 5, (y + h) -5), (0, 255, 255), 3)
        word = prep_image[y:y+h, x:x+w]
        cv2.imwrite(f"./words/{i}_{cntr}.jpg", word)

    ax6.imshow(bbox_image, cmap='gray')
    ax6.title.set_text("Bounding box")